import datetime
import os
import pandas as pd
import requests
import time
from pydantic import BaseModel
from fastmcp import FastMCP, Client
import asyncio

mcp = FastMCP('demo.mcp')

API_BASE_URL = "https://api.financialdatasets.ai"
os.environ["FINANCIAL_DATASETS_API_KEY"] = "80702bee-6402-4bb5-a3c8-829d5207a64d"

class Price(BaseModel):
    open: float
    close: float
    high: float
    low: float
    volume: int
    time: str

class PriceResponse(BaseModel):
    ticker: str
    prices: list[Price]

class FinancialMetrics(BaseModel):
    ticker: str
    report_period: str
    period: str
    currency: str
    market_cap: float | None
    enterprise_value: float | None
    price_to_earnings_ratio: float | None
    price_to_book_ratio: float | None
    price_to_sales_ratio: float | None
    enterprise_value_to_ebitda_ratio: float | None
    enterprise_value_to_revenue_ratio: float | None
    free_cash_flow_yield: float | None
    peg_ratio: float | None
    gross_margin: float | None
    operating_margin: float | None
    net_margin: float | None
    return_on_equity: float | None
    return_on_assets: float | None
    return_on_invested_capital: float | None
    asset_turnover: float | None
    inventory_turnover: float | None
    receivables_turnover: float | None
    days_sales_outstanding: float | None
    operating_cycle: float | None
    working_capital_turnover: float | None
    current_ratio: float | None
    quick_ratio: float | None
    cash_ratio: float | None
    operating_cash_flow_ratio: float | None
    debt_to_equity: float | None
    debt_to_assets: float | None
    interest_coverage: float | None
    revenue_growth: float | None
    earnings_growth: float | None
    book_value_growth: float | None
    earnings_per_share_growth: float | None
    free_cash_flow_growth: float | None
    operating_income_growth: float | None
    ebitda_growth: float | None
    payout_ratio: float | None
    earnings_per_share: float | None
    book_value_per_share: float | None
    free_cash_flow_per_share: float | None

class FinancialMetricsResponse(BaseModel):
    financial_metrics: list[FinancialMetrics]

class LineItem(BaseModel):
    ticker: str
    report_period: str
    period: str
    currency: str
    model_config = {"extra": "allow"}

class LineItemResponse(BaseModel):
    search_results: list[LineItem]

class InsiderTrade(BaseModel):
    ticker: str
    issuer: str | None
    name: str | None
    title: str | None
    is_board_director: bool | None
    transaction_date: str | None
    transaction_shares: float | None
    transaction_price_per_share: float | None
    transaction_value: float | None
    shares_owned_before_transaction: float | None
    shares_owned_after_transaction: float | None
    security_title: str | None
    filing_date: str

class InsiderTradeResponse(BaseModel):
    insider_trades: list[InsiderTrade]

class CompanyNews(BaseModel):
    ticker: str
    title: str
    author: str
    source: str
    date: str
    url: str
    sentiment: str | None = None

class CompanyNewsResponse(BaseModel):
    news: list[CompanyNews]

class CompanyFacts(BaseModel):
    ticker: str
    name: str
    cik: str | None = None
    industry: str | None = None
    sector: str | None = None
    category: str | None = None
    exchange: str | None = None
    is_active: bool | None = None
    listing_date: str | None = None
    location: str | None = None
    market_cap: float | None = None
    number_of_employees: int | None = None
    sec_filings_url: str | None = None
    sic_code: str | None = None
    sic_industry: str | None = None
    sic_sector: str | None = None
    website_url: str | None = None
    weighted_average_shares: int | None = None

class CompanyFactsResponse(BaseModel):
    company_facts: CompanyFacts


def _make_api_request(url: str, headers: dict, method: str = "GET", json_data: dict = None, max_retries: int = 3) -> requests.Response:
    """
    Helper function to make API requests to Financial Datasets API and handle responses.

    Input parameters:
    url: The URL to request.
    headers: Headers to include in the request.
    method: HTTP method (GET or POST, default: GET).
    json_data: JSON data for POST requests (default: None).
    max_retries: Maximum number of retries for rate-limited requests (default: 3).

    Output parameters:
    result: Returns the response object or raises an exception on failure.
    """
    for attempt in range(max_retries + 1):
        if method.upper() == "POST":
            response = requests.post(url, headers=headers, json=json_data)
        else:
            response = requests.get(url, headers=headers)

        if response.status_code == 429 and attempt < max_retries:
            delay = 60 + (30 * attempt)
            print(f"Rate limited (429). Attempt {attempt + 1}/{max_retries + 1}. Waiting {delay}s before retrying...")
            time.sleep(delay)
            continue

        return response


@mcp.tool()
def get_prices(ticker: str, start_date: str, end_date: str, api_key: str = None) -> list[Price]:
    """
    Retrieves daily price data from Financial Datasets API for a given stock ticker.

    Input parameters:
    ticker: The stock ticker symbol (e.g., 'AAPL' for Apple Inc.). This parameter is required.
    start_date: The start date for the price data in YYYY-MM-DD format.
    end_date: The end date for the price data in YYYY-MM-DD format.
    api_key: Optional API key to override environment variable (default: None).

    Output parameters:
    result: Returns a list of Price objects containing price data or an empty list if no data is available.
    """
    headers = {}
    financial_api_key = api_key or os.environ.get("FINANCIAL_DATASETS_API_KEY")
    if financial_api_key:
        headers["X-API-KEY"] = financial_api_key

    url = f"{API_BASE_URL}/prices/?ticker={ticker}&interval=day&interval_multiplier=1&start_date={start_date}&end_date={end_date}"
    response = _make_api_request(url, headers)
    response.raise_for_status()

    price_response = PriceResponse(**response.json())
    return price_response.prices


@mcp.tool()
def get_financial_metrics(ticker: str, end_date: str, period: str = "ttm", limit: int = 10, api_key: str = None) -> list[FinancialMetrics]:
    """
    Retrieves financial metrics from Financial Datasets API for a given stock ticker.

    Input parameters:
    ticker: The stock ticker symbol (e.g., 'AAPL' for Apple Inc.). This parameter is required.
    end_date: The end date for the financial metrics in YYYY-MM-DD format.
    period: The period for the metrics (default: 'ttm' for trailing twelve months).
    limit: Maximum number of records to retrieve (default: 10).
    api_key: Optional API key to override environment variable (default: None).

    Output parameters:
    result: Returns a list of FinancialMetrics objects or an empty list if no data is available.
    """
    headers = {}
    financial_api_key = api_key or os.environ.get("FINANCIAL_DATASETS_API_KEY")
    if financial_api_key:
        headers["X-API-KEY"] = financial_api_key

    url = f"{API_BASE_URL}/financial-metrics/?ticker={ticker}&report_period_lte={end_date}&limit={limit}&period={period}"
    response = _make_api_request(url, headers)
    response.raise_for_status()

    metrics_response = FinancialMetricsResponse(**response.json())
    return metrics_response.financial_metrics


@mcp.tool()
def search_line_items(ticker: str, line_items: list[str], end_date: str, period: str = "ttm", limit: int = 10, api_key: str = None) -> list[LineItem]:
    """
    Retrieves specific financial line items from Financial Datasets API for a given stock ticker.

    Input parameters:
    ticker: The stock ticker symbol (e.g., 'AAPL' for Apple Inc.). This parameter is required.
    line_items: List of financial line items to retrieve (e.g., ['revenue', 'net_income']).
    end_date: The end date for the line items in YYYY-MM-DD format.
    period: The period for the line items (default: 'ttm' for trailing twelve months).
    limit: Maximum number of records to retrieve (default: 10).
    api_key: Optional API key to override environment variable (default: None).

    Output parameters:
    result: Returns a list of LineItem objects or an empty list if no data is available.
    """
    headers = {}
    financial_api_key = api_key or os.environ.get("FINANCIAL_DATASETS_API_KEY")
    if financial_api_key:
        headers["X-API-KEY"] = financial_api_key

    url = f"{API_BASE_URL}/financials/search/line-items"
    body = {
        "tickers": [ticker],
        "line_items": line_items,
        "end_date": end_date,
        "period": period,
        "limit": limit,
    }
    response = _make_api_request(url, headers, method="POST", json_data=body)
    response.raise_for_status()

    response_model = LineItemResponse(**response.json())
    return response_model.search_results[:limit]


@mcp.tool()
def get_insider_trades(ticker: str, end_date: str, start_date: str | None = None, limit: int = 1000, api_key: str = None) -> list[InsiderTrade]:
    """
    Retrieves insider trade data from Financial Datasets API for a given stock ticker.

    Input parameters:
    ticker: The stock ticker symbol (e.g., 'AAPL' for Apple Inc.). This parameter is required.
    end_date: The end date for the insider trades in YYYY-MM-DD format.
    start_date: Optional start date for the insider trades in YYYY-MM-DD format (default: None).
    limit: Maximum number of records to retrieve per request (default: 1000).
    api_key: Optional API key to override environment variable (default: None).

    Output parameters:
    result: Returns a list of InsiderTrade objects or an empty list if no data is available.
    """
    headers = {}
    financial_api_key = api_key or os.environ.get("FINANCIAL_DATASETS_API_KEY")
    if financial_api_key:
        headers["X-API-KEY"] = financial_api_key

    all_trades = []
    current_end_date = end_date

    while True:
        url = f"{API_BASE_URL}/insider-trades/?ticker={ticker}&filing_date_lte={current_end_date}"
        if start_date:
            url += f"&filing_date_gte={start_date}"
        url += f"&limit={limit}"

        response = _make_api_request(url, headers)
        response.raise_for_status()

        response_model = InsiderTradeResponse(**response.json())
        insider_trades = response_model.insider_trades

        if not insider_trades:
            break

        all_trades.extend(insider_trades)

        if not start_date or len(insider_trades) < limit:
            break

        current_end_date = min(trade.filing_date for trade in insider_trades).split("T")[0]

        if current_end_date <= start_date:
            break

    return all_trades


@mcp.tool()
def get_company_news(ticker: str, end_date: str, start_date: str | None = None, limit: int = 1000, api_key: str = None) -> list[CompanyNews]:
    """
    Retrieves company news from Financial Datasets API for a given stock ticker.

    Input parameters:
    ticker: The stock ticker symbol (e.g., 'AAPL' for Apple Inc.). This parameter is required.
    end_date: The end date for the news in YYYY-MM-DD format.
    start_date: Optional start date for the news in YYYY-MM-DD format (default: None).
    limit: Maximum number of records to retrieve per request (default: 1000).
    api_key: Optional API key to override environment variable (default: None).

    Output parameters:
    result: Returns a list of CompanyNews objects or an empty list if no data is available.
    """
    headers = {}
    financial_api_key = api_key or os.environ.get("FINANCIAL_DATASETS_API_KEY")
    if financial_api_key:
        headers["X-API-KEY"] = financial_api_key

    all_news = []
    current_end_date = end_date

    while True:
        url = f"{API_BASE_URL}/news/?ticker={ticker}&end_date={current_end_date}"
        if start_date:
            url += f"&start_date={start_date}"
        url += f"&limit={limit}"

        response = _make_api_request(url, headers)
        response.raise_for_status()

        response_model = CompanyNewsResponse(**response.json())
        company_news = response_model.news

        if not company_news:
            break

        all_news.extend(company_news)

        if not start_date or len(company_news) < limit:
            break

        current_end_date = min(news.date for news in company_news).split("T")[0]

        if current_end_date <= start_date:
            break

    return all_news


@mcp.tool()
def get_market_cap(ticker: str, end_date: str, api_key: str = None) -> float | None:
    """
    Retrieves market capitalization from Financial Datasets API for a given stock ticker.

    Input parameters:
    ticker: The stock ticker symbol (e.g., 'AAPL' for Apple Inc.). This parameter is required.
    end_date: The end date for the market cap in YYYY-MM-DD format.
    api_key: Optional API key to override environment variable (default: None).

    Output parameters:
    result: Returns the market capitalization as a float or None if no data is available.
    """
    headers = {}
    financial_api_key = api_key or os.environ.get("FINANCIAL_DATASETS_API_KEY")
    if financial_api_key:
        headers["X-API-KEY"] = financial_api_key

    if end_date == datetime.datetime.now().strftime("%Y-%m-%d"):
        url = f"{API_BASE_URL}/company/facts/?ticker={ticker}"
        response = _make_api_request(url, headers)
        if response.status_code != 200:
            print(f"Error fetching company facts: {ticker} - {response.status_code}")
            return None

        response_model = CompanyFactsResponse(**response.json())
        return response_model.company_facts.market_cap

    financial_metrics = get_financial_metrics(ticker, end_date, api_key=api_key)
    if not financial_metrics:
        return None

    return financial_metrics[0].market_cap

GET_PRICES_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_prices",
        "description": "Retrieves daily price data from Financial Datasets API for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple Inc.)."
                },
                "start_date": {
                    "type": "string",
                    "description": "The start date for the price data in YYYY-MM-DD format."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the price data in YYYY-MM-DD format."
                },
                "api_key": {
                    "type": "string",
                    "description": "Optional API key to override environment variable.",
                    "default": None
                }
            },
            "required": ["ticker", "start_date", "end_date"]
        }
    }
}

GET_FINANCIAL_METRICS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_financial_metrics",
        "description": "Retrieves financial metrics from Financial Datasets API for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple Inc.)."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the financial metrics in YYYY-MM-DD format."
                },
                "period": {
                    "type": "string",
                    "description": "The period for the metrics (e.g., 'ttm' for trailing twelve months).",
                    "default": "ttm"
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of records to retrieve.",
                    "default": 10
                },
                "api_key": {
                    "type": "string",
                    "description": "Optional API key to override environment variable.",
                    "default": None
                }
            },
            "required": ["ticker", "end_date"]
        }
    }
}

SEARCH_LIN_ITEMS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "search_line_items",
        "description": "Retrieves specific financial line items from Financial Datasets API for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple Inc.)."
                },
                "line_items": {
                    "type": "array",
                    "items": {
                        "type": "string"
                    },
                    "description": "List of financial line items to retrieve (e.g., ['revenue', 'net_income'])."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the line items in YYYY-MM-DD format."
                },
                "period": {
                    "type": "string",
                    "description": "The period for the line items (e.g., 'ttm' for trailing twelve months).",
                    "default": "ttm"
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of records to retrieve.",
                    "default": 10
                },
                "api_key": {
                    "type": "string",
                    "description": "Optional API key to override environment variable.",
                    "default": None
                }
            },
            "required": ["ticker", "line_items", "end_date"]
        }
    }
}

GET_INSIDER_TRADES_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_insider_trades",
        "description": "Retrieves insider trade data from Financial Datasets API for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple Inc.)."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the insider trades in YYYY-MM-DD format."
                },
                "start_date": {
                    "type": "string",
                    "description": "Optional start date for the insider trades in YYYY-MM-DD format.",
                    "default": None
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of records to retrieve per request.",
                    "default": 1000
                },
                "api_key": {
                    "type": "string",
                    "description": "Optional API key to override environment variable.",
                    "default": None
                }
            },
            "required": ["ticker", "end_date"]
        }
    }
}

GET_COMPANY_NEWS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_company_news",
        "description": "Retrieves company news from Financial Datasets API for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple Inc.)."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the news in YYYY-MM-DD format."
                },
                "start_date": {
                    "type": "string",
                    "description": "Optional start date for the news in YYYY-MM-DD format.",
                    "default": None
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of records to retrieve per request.",
                    "default": 1000
                },
                "api_key": {
                    "type": "string",
                    "description": "Optional API key to override environment variable.",
                    "default": None
                }
            },
            "required": ["ticker", "end_date"]
        }
    }
}

GET_MARKET_CAP = {
    "type": "function",
    "function": {
        "name": "get_market_cap",
        "description": "Retrieves market capitalization from Financial Datasets API for a given stock ticker.",
        "parameters": {
            "type": "object",
            "properties": {
                "ticker": {
                    "type": "string",
                    "description": "The stock ticker symbol (e.g., 'AAPL' for Apple Inc.)."
                },
                "end_date": {
                    "type": "string",
                    "description": "The end date for the market cap in YYYY-MM-DD format."
                },
                "api_key": {
                    "type": "string",
                    "description": "Optional API key to override environment variable.",
                    "default": None
                }
            },
            "required": ["ticker", "end_date"]
        }
    }
}


async def main():
    client = Client(mcp)
    async with client:
        # View available tools
        tools = await client.list_tools()
        print('Available tools:', tools)
        
        # Call the tools (parameter names must match)
        prices_result = await client.call_tool('get_prices', {'ticker': 'AAPL', 'start_date': '2025-01-01', 'end_date': '2025-10-15'})
        print('Prices Result:', prices_result)
        
        financial_metrics_result = await client.call_tool('get_financial_metrics', {'ticker': 'AAPL', 'end_date': '2025-10-15'})
        print('Financial Metrics Result:', financial_metrics_result)
        
        line_items_result = await client.call_tool('search_line_items', {'ticker': 'AAPL', 'line_items': ['revenue', 'net_income'], 'end_date': '2025-10-15'})
        print('Line Items Result:', line_items_result)
        
        insider_trades_result = await client.call_tool('get_insider_trades', {'ticker': 'AAPL', 'end_date': '2025-10-15', 'start_date': '2025-01-01'})
        print('Insider Trades Result:', insider_trades_result)
        
        company_news_result = await client.call_tool('get_company_news', {'ticker': 'AAPL', 'end_date': '2025-10-15', 'start_date': '2025-01-01'})
        print('Company News Result:', company_news_result)
        
        market_cap_result = await client.call_tool('get_market_cap', {'ticker': 'AAPL', 'end_date': '2025-10-15'})
        print('Market Cap Result:', market_cap_result)

if __name__ == '__main__':
    asyncio.run(main())